import tensorflow as tf
import os
import matplotlib.pyplot as plt
import numpy as np
import random

#定义读入图片文件函数
def readimg(file):
    image = plt.imread(file)
    image = image / 255. #归一化，缩放
    return image

images = []   #存储所有图片文件的路径
labels = []   #存储所有图片的独热之后的y
path = r'data3\data3'
for f in os.listdir(path):  #循环列出所有文件名:如0_6.jpg
    img = readimg(path + '\\' + f) #比如,f: 0_6.jpg, f[0]=0表示第一个字符, 表示分类类别
    images.append(img)
    ones = np.eye(3)   #单位矩阵
    labels.append(ones[int(f[0])])  # 独热编码, f: 0_6.jpg, f[0]=0表示第一个字符, 表示分类类别0

#list转化为np.array
images_data = np.array(images)
labels = np.array(labels)
print(images_data.shape)

# 洗牌乱序: 打乱顺序
np.random.seed(3)
m = images_data.shape[0]  #样本个数
order = np.random.permutation(m)  #产生长度为m的随机序列
imgArr = images_data[order]
Y_one_hot = labels[order]
total = imgArr.shape[0]
train_test = 0.9

g_b=0
# 自己实现next_batch函数，每次返回一批数据
def next_batch(size):
    global g_b
    xb = imgArr[g_b:g_b+size]
    yb = Y_one_hot[g_b:g_b+size]
    g_b = g_b + size
    return xb,yb

# 参数
learning_rate = 0.001 # 学习率
training_epochs = 40  # 训练总周期
batch_size = 10 # 训练每批样本数

#定义占位符
X = tf.placeholder(tf.float32, [None, 32, 32, 3])
Y = tf.placeholder(tf.float32, [None, 3])  # 独热编码

# 第1层卷积，输入图片数据(?, 32, 32, 3)
W1 = tf.Variable(tf.random_normal([3, 3, 3, 32]))  #卷积核3x3，输入通道3，输出通道32
L1 = tf.nn.conv2d(X, W1, strides=[1, 1, 1, 1], padding='SAME') #卷积输出 （?, 32, 32, 32)
L1 = tf.nn.relu(L1)
L1 = tf.nn.max_pool(L1, ksize=[1, 2, 2, 1],strides=[1, 2, 2, 1], padding='SAME') #池化输出 (?, 16, 16, 32)

# 第2层卷积，输入图片数据(?, 16, 16, 32)
W2 = tf.Variable(tf.random_normal([3, 3, 32, 64], stddev=0.01)) #卷积核3x3，输入通道32，输出通道64
L2 = tf.nn.conv2d(L1, W2, strides=[1, 1, 1, 1], padding='SAME') #卷积输出  (?, 16, 16, 64)
L2 = tf.nn.relu(L2)
L2 = tf.nn.max_pool(L2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME') #池化输出 (?, 8, 8, 64)

# 展开：变成一维向量
dim = L2.get_shape()[1].value * L2.get_shape()[2].value * L2.get_shape()[3].value
L2_flat = tf.reshape(L2,[-1, dim])

# 全连接  -> 3 outputs
W3 = tf.get_variable("W3", shape=[dim, 3], initializer=tf.contrib.layers.xavier_initializer())
b = tf.Variable(tf.random_normal([3]))
logits = tf.matmul(L2_flat, W3) + b

#计算准确率
correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(Y, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

#代价或损失函数
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=Y))
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost) # 优化器

# 创建会话
sess = tf.Session()
sess.run(tf.global_variables_initializer()) #全局变量初始化
# 迭代训练
print('开始学习...')
for epoch in range(training_epochs):
    avg_cost = 0
    total_batch = int(total * train_test / batch_size)  # 计算总批次(训练集)
    g_b = 0 #必须复位
    for i in range(total_batch):
        batch_xs, batch_ys = next_batch(batch_size)
        feed_dict = {X: batch_xs, Y: batch_ys}
        c, _ = sess.run([cost, optimizer], feed_dict=feed_dict)
        avg_cost += c / total_batch
    acc = sess.run(accuracy, feed_dict={X: imgArr[int(total * train_test):], Y: Y_one_hot[int(total * train_test):]})
    print('Epoch:', (epoch + 1), 'cost =', avg_cost, 'acc =', acc)
print('学习完成')

# 测试模型检查准确率
print('Accuracy:', sess.run(accuracy, feed_dict={X: imgArr[int(total * train_test):], Y: Y_one_hot[int(total * train_test):]}))

# 在测试集中随机抽一个样本进行测试
r = int(total * train_test) + random.randint(0, int(total * (1 - train_test)) - 1)
print("Label: ", sess.run(tf.argmax(Y_one_hot[r:r + 1], 1)))
print("Prediction: ", sess.run(tf.argmax(logits, 1), feed_dict={X: imgArr[r:r + 1]}))
